DisaggKD: Efficient Knowledge Distillation through Teacher-Student Disaggregation and a Fused Loss
DisaggKD：通过师生分离与融合损失实现高效知识蒸馏
🧭 知识蒸馏训练效率优化
The University of Tokyo · RIKEN AIP · SB Intuitions Corp.
提出教师-学生分离的KD框架，通过fused loss和prefetch ring突破显存限制，匹配同布局损失并提速
📌 工作参考价值（4 点）
- 在资源受限下做LLM蒸馏（例如蒸馏通用大模型到电商领域小模型用于query理解或排序），采用教师-学生分离部署，教师只前向，学生负责训练，避免同一GPU同时存放两套权重和优化器状态。
- 实现KD loss时，若词表过大，不必显式构造完整logits张量，可通过行窗口和词表分片（vocabulary shards）的融合kernel精确计算loss，既省显存又保持数值一致，适合大词表场景（如搜索词预测、文案生成）。
- 在跨节点流水线并行训练中引入prefetch ring协调生产者和训练者，可掩盖通信延迟，提高吞吐，对需要跨节点数据转移的架构有借鉴意义。
- 使用闭式内存和流量模型提前判断不同布局（colocated vs 分离）的可行性，避免OOM，可在设计训练方案时快速评估，节省试错成本。